import streamlit as st
import os
from dotenv import load_dotenv
load_dotenv()  # 加载.env文件中的环境变量

import sqlite3
import datetime
from dotenv import load_dotenv
from langchain_community.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.agents import initialize_agent, AgentType
from langchain.tools import Tool
from sqlalchemy import create_engine, text
from contextlib import contextmanager
from main import directory_agent, database_agent

# 加载环境变量
load_dotenv()

# 初始化页面配置
st.set_page_config(page_title="笔记管理系统", page_icon="📝", layout="wide")

# 数据库配置
DB_PATH = "notes_database.db"

# 初始化数据库连接
@contextmanager
def db_connection():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    try:
        yield cursor
        conn.commit()
    except Exception as e:
        conn.rollback()
        st.error(f"数据库错误: {str(e)}")
    finally:
        conn.close()

# 初始化数据库表结构
def init_database():
    with db_connection() as cursor:
        # 伪日志表 - 记录所有SQL操作
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS sql_operations_log (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            sql_command TEXT,
            success BOOLEAN
        )
        ''')

        # 文件表 - 存储文件信息和标签
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS files (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            file_path TEXT UNIQUE,
            file_name TEXT,
            folder_path TEXT,
            created_at DATETIME,
            updated_at DATETIME,
            tags TEXT,
            summary TEXT
        )
        ''')

        # 模型调用记录表 - 记录外部模型调用
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS model_calls (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            user_query TEXT,
            model_response TEXT,
            sql_operations TEXT
        )
        ''')

# 记录SQL操作到日志表
def log_sql_operation(sql_command, success=True):
    with db_connection() as cursor:
        cursor.execute(
            "INSERT INTO sql_operations_log (timestamp, sql_command, success) VALUES (?, ?, ?)",
            (datetime.datetime.now(), sql_command, success)
        )

# 使用LLM提取文件分类关键词
def get_categories_from_llm(content, llm):
    prompt = PromptTemplate(
        input_variables=["content"],
        template="分析以下文本内容，返回1-3个适合作为文件夹名称的关键词，用逗号分隔。关键词应简洁且能代表文本主题。文本内容: {content}"
    )
    chain = LLMChain(llm=llm, prompt=prompt)
    try:
        response = chain.run(content).strip()
        # 解析响应，处理可能的格式问题
        categories = [cat.strip() for cat in response.split(',') if cat.strip()]
        return categories[:3] if categories else ["未分类"]
    except Exception as e:
        st.warning(f"分类失败，使用默认分类: {str(e)}")
        return ["未分类"]

# 记录模型调用
def log_model_call(user_query, model_response, sql_operations):
    with db_connection() as cursor:
        cursor.execute(
            "INSERT INTO model_calls (timestamp, user_query, model_response, sql_operations) VALUES (?, ?, ?, ?)",
            (datetime.datetime.now(), user_query, model_response, sql_operations)
        )

# 生成文件摘要
def generate_summary(content):
    if len(content) < 100:
        return content
    return content[:100] + "..."

# 组织文件到合理结构的文件夹
def organize_files(input_folder):
    if not os.path.exists(input_folder):
        st.error(f"输入文件夹不存在: {input_folder}")
        return

    # 创建主输出文件夹
    output_folder = os.path.join(os.path.dirname(input_folder), "organized_notes")
    os.makedirs(output_folder, exist_ok=True)

    sql_commands = []

    # 遍历输入文件夹中的所有MD文件
    for root, dirs, files in os.walk(input_folder):
        for file in files:
            if file.endswith(".md"):
                file_path = os.path.join(root, file)
                try:
                    # 读取文件内容
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()

                    # 提取关键词作为文件夹名称
                    keywords = get_categories_from_llm(content, get_base_llm())
                    if not keywords:
                        keywords = ["未分类"]

                    # 创建文件夹结构
                    folder_path = os.path.join(output_folder, *keywords)
                    os.makedirs(folder_path, exist_ok=True)

                    # 复制文件到新位置
                    new_file_path = os.path.join(folder_path, file)
                    with open(new_file_path, 'w', encoding='utf-8') as f:
                        f.write(content)

                    # 生成摘要文件
                    summary = generate_summary(content)
                    with open(os.path.join(folder_path, "README.md"), 'w', encoding='utf-8') as f:
                        f.write(f"# {folder_path.split(os.sep)[-1]} 文件夹\n\n")
                        f.write(f"## 文件夹信息\n\n")
                        f.write(f"- 创建时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                        f.write(f"- 包含文件: {len(os.listdir(folder_path)) - 1}\n\n")
                        f.write(f"## 内容摘要\n\n{summary}")

                    # 记录文件信息到数据库
                    tags = ",".join(keywords)
                    # 使用参数化查询避免SQL注入和语法错误
                    sql = "INSERT OR IGNORE INTO files (file_path, file_name, folder_path, created_at, updated_at, tags, summary) VALUES (?, ?, ?, ?, ?, ?, ?)"
                    with db_connection() as cursor:
                        cursor.execute(sql, (new_file_path, file, folder_path, datetime.datetime.now(), datetime.datetime.now(), tags, summary))
                    sql_commands.append(sql)
                    log_sql_operation(sql)

                except Exception as e:
                    st.error(f"处理文件 {file} 时出错: {str(e)}")

    st.success(f"文件组织完成！已保存到: {output_folder}")
    return sql_commands

# 初始化基础LLM模型
def get_base_llm():
    return ChatOpenAI(
        temperature=0,
        model="deepseek-chat",
        base_url="https://api.deepseek.com/v1",
        api_key=os.getenv("DEEPSEEK_API_KEY")
    )

# 初始化主agent，整合两个MCP代理
def init_main_agent():
    llm = get_base_llm()
    tools = [
        Tool(
            name="DirectoryAgent",
            func=lambda query: directory_agent.run(query),
            description="用于管理文件系统目录的工具，处理文件列表、创建、读取等操作"
        ),
        Tool(
            name="DatabaseAgent",
            func=lambda query: database_agent.run(query),
            description="用于操作SQLite数据库的工具，处理查询和修改操作"
        )
    ]
    # 系统提示词，指导agent如何使用工具
    system_message = '''你是一个笔记管理管家，专门负责帮助用户管理笔记文件和数据库。当被问及身份时，必须回答：'我是你的笔记管理管家，可帮你管理文件系统和数据库。'需要根据用户需求选择合适的工具：
1. 当需要管理文件系统（如创建/读取文件、列出目录）时，使用DirectoryAgent
2. 当需要操作数据库（如查询/修改笔记内容）时，使用DatabaseAgent
3. 先分析用户问题，再决定调用哪个工具，不需要同时调用两个工具
4. 直接返回工具的执行结果，不需要额外解释'''

    return initialize_agent(
        tools, llm, agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True, agent_kwargs={"system_message": system_message}
    )

# 处理用户查询
def process_query(query):
    # 使用主agent处理查询，自动路由到相应的MCP代理
    main_agent = init_main_agent()
    return main_agent.run(query)

# 数据库查询函数
def query_database(query):
    engine = create_engine(f"sqlite:///{DB_PATH}")
    try:
        with engine.connect() as conn:
            result = conn.execute(text(query))
            columns = result.keys()
            return [dict(zip(columns, row)) for row in result.fetchall()]
    except Exception as e:
        st.error(f"查询错误: {str(e)}")
        return []

# 主应用
def main():
    st.title("📝 智能笔记管理系统")
    st.subheader("基于LangChain的本地文件与数据库管理")

    # 初始化数据库
    init_database()

    # 初始化LLM
    llm = get_base_llm()

    # 创建标签页
    tab1, tab2 = st.tabs(["助手对话", "数据库查看"])

    # 侧边栏 - 文件组织
    with st.sidebar:
        st.header("文件组织")
        input_folder = st.text_input("输入文件夹路径")
        if st.button("组织文件"):
            if input_folder:
                if not llm:
                    st.error("请先设置DEEPSEEK_API_KEY以使用AI分类功能")
                else:
                    with st.spinner("正在使用AI分析文件内容并组织..."):
                        sql_commands = organize_files(input_folder)
                        if sql_commands:
                            st.session_state["last_sql_commands"] = ",".join(sql_commands)

    with tab1:
        # 聊天交互界面
        st.header("与助手对话")

        # 初始化聊天历史
        if "chat_history" not in st.session_state:
            st.session_state["chat_history"] = []
        if "last_sql_commands" not in st.session_state:
            st.session_state["last_sql_commands"] = ""

        # 显示聊天历史
        for message in st.session_state["chat_history"]:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

    # 用户输入 - 必须放在标签页外部
    user_query = st.chat_input("请输入您的查询...")
    if user_query:
        # 显示用户消息
        with st.chat_message("user"):
            st.markdown(user_query)
        st.session_state["chat_history"].append({"role": "user", "content": user_query})

        # 处理查询
        with st.spinner("思考中..."):
            response = process_query(user_query)

            # 记录模型调用
            log_model_call(user_query, response, st.session_state["last_sql_commands"])
            st.session_state["last_sql_commands"] = ""

            # 显示助手回复
            with st.chat_message("assistant"):
                st.markdown(response)
            st.session_state["chat_history"].append({"role": "assistant", "content": response})

    with tab2:
        # 数据库查看界面
        st.header("数据库内容查看")
        table_option = st.selectbox(
            "选择要查看的表",
            ("files", "model_calls", "sql_operations_log")
        )

        if st.button("加载数据"):
            with st.spinner(f"正在加载{table_option}表数据..."):
                if table_option == "files":
                    data = query_database("SELECT id, file_name, folder_path, tags, created_at FROM files ORDER BY created_at DESC")
                elif table_option == "model_calls":
                    data = query_database("SELECT id, timestamp, user_query, model_response FROM model_calls ORDER BY timestamp DESC LIMIT 50")
                else:
                    data = query_database("""
                        SELECT 
                            l.id, l.timestamp, l.sql_command, l.success, 
                            m.user_query, m.model_response
                        FROM sql_operations_log l
                        LEFT JOIN model_calls m ON l.model_call_id = m.id
                        ORDER BY l.timestamp DESC LIMIT 50
                    """)

                if data:
                    st.dataframe(data, use_container_width=True)
                else:
                    st.info("该表暂无数据")

if __name__ == "__main__":
    main()